import numpy as np
from collections import deque

def robust_solition_distribution(K, C = 0.1, delta = 0.5):
    """
    计算鲁棒孤波分布
    
    参数：
    K -- 数据块的数量
    C -- 鲁棒系数，整数参数
    delta -- 偏差系数，失败概率
    
    返回：
    mu -- 鲁棒孤波分布
    """
    # 理想孤波分布
    rho = np.zeros(K)
    rho[0] = 1 / K
    for i in range(1, K):
        rho[i] = 1 / (i * ( i + 1))
        
    # 鲁棒系数
    R = C * np.log(K / delta) * np.sqrt(K)
    if R <= 0:
        R = 1   # 鲁棒系数C太小时，R可能为负数，取1代替

    # 鲁棒孤波分布的修正项
    tau = np.zeros(K)
    for i in range(K):
        if i <= (K / R) - 1:
            tau[i] = R / (K * (i + 1))      # i是从0开始的，在计算时要+1
        elif i == int(K / R) - 1:
            tau[i] = (R * np.log(R / delta)) / K
        else:
            tau[i] = 0
    
    # 归一化因子
    Z = np.sum(rho + tau)

    # 鲁棒孤波分布
    mu = (rho + tau) / Z
    return mu

def generate_degree(mu):
    """
    根据鲁棒孤波分布生成随机度数
    
    参数：
    mu -- 鲁棒孤波分布概率

    返回：
    d -- 随机度数
    """
    return np.random.choice(np.arange(1, len(mu) + 1), p = mu)

def lt_generate_matrix(K):
    """
    LT编码生成矩阵 G, 1 X K 矩阵

    参数：
    K -- 数据块的数量, 即data的行数
    
    返回：
    d -- 随机度数
    G -- LT编码矩阵
    """
    mu = robust_solition_distribution(K)
    G = np.zeros((1, K), dtype = int)
    d = generate_degree(mu)
    indices = np.random.choice(K, size = d, replace = False)
    G[0, indices] = 1   # 对被选中的数据包对应行列置1
    return d, G

def lt_encode(data, K, N):
    """
    LT编码 - 单个LT包的编码过程
    
    参数：
    data -- 数据块, K X N 二维数组
    K -- 数据块的数量
    N -- 数据块的长度
    
    返回：
    d -- 随机度数
    LT_package -- 编码后的LT包, 1 X N 矩阵
    G -- 生成矩阵的行, 即挑选的数据包(1 X K 矩阵)
    """
    d, G = lt_generate_matrix(K)
    LT_package = np.zeros((1, N), dtype = int)
    # 根据G中为1的列挑选data中对应的数据包
    indices = np.where(G[0, :] == 1)[0]     # 获取被选中的列的索引
    if len(indices) > 0:
        # 对被选中的数据包进行异或操作，生成LT包
        LT_package = np.bitwise_xor.reduce(data[indices], axis = 0)
        
    return d, LT_package, G[0]

def lt_decode_gaussian(received_LT_packages, received_G, K, N):
    """
    LT的高斯消元解码算法

    参数：
    received_LT_packages -- 接收的编码包, M X N 二维数组
    received_G -- 生成矩阵, M X K 矩阵, M是编码包的数量
    K -- 原始数据包的数量
    N -- 数据包的长度

    返回:
    decoded_data -- 解码后的原始数据, K X N 二维数组
    success_flag -- 是否完全解码的标志
    """
    decoded_data = np.zeros((K, N), dtype = int)
    success_flag = True
    M = received_G.shape[0]

    # 检查最小解码条件
    if M < K:
        print(f"!!! Warning: M < K, decoding may fail")
        success_flag = False
        return decoded_data, success_flag

    # 显示转换为整型矩阵
    A = received_G.astype(int)
    B = received_LT_packages.astype(int)

    for col in range(N):
        
        augmented = np.hstack((A, B[:, col].reshape(-1, 1)))
        rank = 0
        rows, cols = augmented.shape

        # 高斯消元（仅处理前K列）
        for pivot_col in range(K):
            # 寻找主元行(当前列第一个1所在的行)
            pivot_row = None
            for r in range(rank, rows):
                if augmented[r, pivot_col] == 1:
                    pivot_row = r
                    break
            # 该列无主元, 跳过
            if pivot_row is None:
                continue

            # 交换主元行与当前行
            augmented[[rank, pivot_row]] = augmented[[pivot_row, rank]]

            # 消去其他行中该列的值
            for r in range(rows):
                if r != rank and augmented[r, pivot_col] == 1:
                    augmented[r] ^= augmented[rank]
            rank += 1

        # 检查是否存在解 （0 = 1的矛盾）
        for r in range(rank, rows):
            if augmented[r, -1] != 0:
                print(f"!!! Warning: Column {col} has contradictory equations, decoding failed")
                success_flag = False
                decoded_data[:, col] = 0
                break

        # 检查秩是否足够解方程
        if rank < K:
            print(f"!!! Warning: Column {col} has insufficient rank({rank} < {K}), decoding failed")
            success_flag = False
            decoded_data[:, col] = 0
            continue

        # 回代，求解
        x = np.zeros(K, dtype = int)
        for r in reversed(range(rank)):
            # 找到主元所在列
            cols = np.where(augmented[r, :K] == 1)[0]
            if len(cols) == 0:
                continue    # 该行全0, 跳过
            col_idx = cols[0]
            x[col_idx] = augmented[r, -1]
            # 消去上方行中该列的值
            for i in range(r):
                if augmented[i, col_idx] == 1:
                    augmented[i, -1] ^= x[col_idx]
                    # augmented[i, col_idx] = 0    # 仅为了清晰，实际在GF(2)中无需此操作
        decoded_data[:, col] = x

    return decoded_data, success_flag

def lt_decode_bp(received_LT_packages, received_G, K, N):
    """
    LT的BP解码算法: 

    参数：
    received_LT_packages -- 接收的编码包, M X N 二维数组
    received_G -- 生成矩阵, M X K 矩阵, M是编码包的数量
    K -- 原始数据包的数量
    N -- 数据包的长度

    返回:
    decoded_data -- 解码后的原始数据, K X N 二维数组
    success_flag -- 是否完全解码的标志
    """
    decoded_data = np.zeros((K, N), dtype = int)
    success_flag = True
    M = received_G.shape[0]
    # 检查最小解码条件
    if M < K:
        print(f"!!! Warning: M < K, decoding may fail")
        success_flag = False
        return decoded_data, success_flag
    # 将生成矩阵转换为稀疏表示(提高处理效率)
    sparse_G = [set(np.where(row == 1)[0]) for row in received_G]
    # 初始化数据结构(整个数据包级别)
    decoded = np.zeros(K, dtype = bool)
    remaining_equations = set(range(M))
    equation_degrees = np.array([len(indices) for indices in sparse_G], dtype = int)
    data_to_equations = [set() for _ in range(K)]
    #建立初始关联关系(形成图连接)
    for eq_idx, data_indices in enumerate(sparse_G):
        for d in data_indices:
            data_to_equations[d].add(eq_idx)
    # 初始化队列(度数为1的方程)
    queue = deque([eq for eq in remaining_equations if equation_degrees[eq] == 1])

    while queue:
        eq_idx = queue.popleft()
        # 跳过已处理或度数变化的方程
        if eq_idx not in remaining_equations or equation_degrees[eq_idx] != 1:
            continue
        # 提取关联数据包
        try:
            data_idx = next(iter(sparse_G[eq_idx]))
        except StopIteration:
            remaining_equations.discard(eq_idx)
            continue    # 无关联数据包的情况
        # 如果数据包已解码则跳过
        if decoded[data_idx]:
            continue
        # 解码数据包(所有列同时处理)
        decoded_data[data_idx] = received_LT_packages[eq_idx]
        decoded[data_idx] = True
        # 更新关联关系(并行处理所有列)
        related_equations = data_to_equations[data_idx].copy()
        for other_eq in related_equations:
            if other_eq not in remaining_equations:
                continue
            # 更新方程值(异或操作)
            received_LT_packages[other_eq] = np.bitwise_xor(received_LT_packages[other_eq], decoded_data[data_idx])
            # 更新生成矩阵关系
            sparse_G[other_eq].discard(data_idx)
            equation_degrees[other_eq] = len(sparse_G[other_eq])
            # 更新双向映射关系
            data_to_equations[data_idx].discard(other_eq)
            # 维护队列状态
            if equation_degrees[other_eq] == 1:
                queue.append(other_eq)
            elif equation_degrees[other_eq] == 0:
                remaining_equations.discard(other_eq)
        # 标记方程已处理
        remaining_equations.discard(eq_idx)
        # equation_degrees[eq_idx] = 0
    # 最终完整性检查
    success_flag = np.all(decoded)

    return decoded_data, success_flag



def binary_symmetric_channel(data, ber = 0.05):
    """
    二进制对称信道(BSC): 每个比特有一定的概率发生翻转
    在此信道下有一个突出问题: 比特翻转后，接收方无法判断是否是误码，导致解码失败

    参数:
    data -- 数据块, K X N 二维数组
    ber -- 误比特率

    返回:
    corrupted_data -- 带有BSC的错误数据块, K X N 二维数组
    """
    # 生成错误掩码, 1表示发生错误, 0表示无错误
    error_mask = np.random.choice([0, 1], size = data.shape, p = [1 - ber, ber])    # data.shape = (K, N)
    # 异或操作产生错误数据, 异或后1会发生翻转, 0保持不变
    corrupted_data = np.bitwise_xor(data, error_mask)
    return corrupted_data, error_mask

def erasure_channel(data, loss_prob = 0.1):
    """
    擦除信道(EC): 随机丢弃一部分数据块

    参数:
    data -- 数据块, K X N 二维数组
    loss_prob -- 丢失概率

    返回:
    corrupted_data -- 发生丢失后的原数据块, K X N 二维数组
    """
    # 生成丢失掩码, 1表示未丢失, 0表示丢失
    loss_mask = np.random.choice([1, 0], size = data.shape, p = [1 - loss_prob, loss_prob])
    # 乘法操作丢失数据
    corrupted_data =  data[loss_mask]
    return corrupted_data, loss_mask

def simulate_lt_transmission(K, N, decode_type = 'gaussian', channel_type = 'BSC', error_param = 0.05):
    """
    模拟LT编码-传输-解码过程

    参数:
    K -- 数据块的数量
    N -- 数据块的长度
    decode_type -- 解码算法类型, 'gaussian' or 'bp'
    channel_type -- 信道类型, 'BSC' or 'EC'
    error_param -- 错误概率, 0.1 or 0.05

    返回:
    decoded_data -- 解码后的原始数据, K X N 二维数组
    success_flag -- 是否完全解码的标志
    """
    # 1.生成原始数据
    original_data = np.random.randint(0, 2, size = (K, N))

    # 2.LT编码
    LT_packages = np.zeros((int(1.5 * K), N), dtype = int)
    generate_matrix = np.zeros((int(1.5 * K), K), dtype = int)
    d_sum = np.zeros((1, int(1.5 * K)), dtype = int)
    for i in range(int(1.5 * K)):
        d, LT_package, G_row = lt_encode(original_data, K, N)
        LT_packages[i, :] = LT_package
        generate_matrix[i, :] = G_row
        d_sum[0, i] = d
    d_proportion = np.sum(d_sum == 1, axis = 1) / int(1.5 * K)

    # 3.信道干扰
    if channel_type == 'BSC':
        corrupted_data, error_mask = binary_symmetric_channel(LT_packages, ber = error_param)
    elif channel_type == 'EC':
        corrupted_data, loss_mask = erasure_channel(LT_packages, loss_prob = error_param)
    else:
        print(f"!!! Warning: Invalid channel type {channel_type}, use BSC instead")
        corrupted_data, error_mask = binary_symmetric_channel(LT_packages, ber = error_param)

    # 4.解码
    if decode_type == 'gaussian':
        decoded_data, success_flag = lt_decode_gaussian(LT_packages, generate_matrix, K, N)
    elif decode_type == 'bp':
        decoded_data, success_flag = lt_decode_bp(LT_packages, generate_matrix, K, N)
    else:
        print(f"!!! Warning: Invalid decode type {decode_type}, use bp instead")
        decoded_data, success_flag = lt_decode_bp(LT_packages, generate_matrix, K, N)
    
    # 5.统计解码成功率
    success_proportion = np.sum(np.all(decoded_data == original_data, axis = 1)) / K

    return original_data, decoded_data, success_flag, success_proportion, d_proportion



if __name__ == '__main__':
    # 测试参数
    K = 100
    N = 128
    trials = 50
    success_count = 0

    # 运行多次测试
    for i in range(1, trials + 1):
    # 模拟LT完整传输过程
        original_data, decoded_data, success_flag, success_proportion, d_proportion = \
            simulate_lt_transmission(
                K,
                N,
                decode_type = 'bp',   # 'gaussian' or 'bp'
                channel_type = 'BSC',   # 'BSC'0.05 or 'EC'0.1
                error_param = 0.05
                )
        # 打印结果
        print(f"The {i}th trial:")
        print(f"The proportion of encoded packages with degree 1: {d_proportion}")
        print("Original data:\n", original_data)
        print("Decoded data:\n", decoded_data)
        print(f"Success flag:{success_flag}")
        print(f"Success proportion:{success_proportion}" )
        print("-" * 50)
        print()
        if success_flag:
            success_count += 1
    print(f"The success rate is {success_count / trials:.2f}")

  






"""
说明:此代码为单纯LT码 编码-传输-解码过程的模拟，并未涉及信道干扰、错误处理等过程。
包含以下功能：
1. 生成原始数据
2. LT编码
# 3. 信道干扰(未实现)
4. 解码
5. 统计解码成功率


截止2025.2.17，LT高斯消元解码和BP解码算法都成功实现，且成功率较高，但存在以下问题

情况1: 高斯消元算法由于时间复杂度较高，在大数据量时运行较慢，BP算法不存在此情况
情况2: LT编译码可实现，但是经过信道会产生较大误差，导致解码失败，还未解决

"""



